import subprocess

class GPUDetector:
    def __init__(self):
        self.gpu_info = self._detect_gpu()

    def _detect_gpu(self):
        """
        调用 lspci 获取显卡信息
        返回格式: {"vendor": "NVIDIA", "model": "GeForce GTX 1060"}
        """
        try:
            output = subprocess.check_output(["lspci", "-v"], text=True)
            for line in output.splitlines():
                if "VGA" in line or "Display controller" in line:
                    _, rest = line.split(":", 1)
                    vendor, model = rest.strip().split(" ", 1)
                    vendor = vendor.lower()
                    if "nvidia" in vendor:
                        return {"vendor": "NVIDIA", "model": model}
                    elif "intel" in vendor:
                        return {"vendor": "Intel", "model": model}
                    elif "amd" in vendor or "ati" in vendor:
                        return {"vendor": "AMD", "model": model}
            return {"vendor": "Unknown", "model": ""}
        except Exception as e:
            print("检测显卡失败:", e)
            return {"vendor": "Unknown", "model": ""}

    def recommend_kernel_params(self):
        """
        根据显卡推荐内核参数
        """
        gpu_vendor = self.gpu_info["vendor"]
        recommendations = {
            "NVIDIA": ["nomodeset"],
            "AMD": ["radeon.modeset=0"],
            "Intel": ["i915.modeset=0"],
        }
        return recommendations.get(gpu_vendor, [])